"""
索引存储管理模块

负责向量索引的存储、加载和管理。
"""

import os
import time
from typing import Optional, List, Dict, Any

from llama_index.core import VectorStoreIndex, Document
from llama_index.core import StorageContext
from llama_index.core import load_index_from_storage

from .storage_base import StorageBase


class IndexStorage(StorageBase):
    """
    向量索引存储管理器

    职责：
    1. 管理向量索引的存储和加载
    2. 提供索引生命周期管理
    3. 监控索引存储状态
    4. 提供索引备份和恢复功能
    """

    def __init__(self, index_dir: Optional[str] = None):
        """
        初始化索引存储管理器

        Args:
            index_dir: 索引存储目录
        """
        if index_dir is None:
            from ..config import config
            index_dir = config.index_dir

        super().__init__(index_dir)
        self._index: Optional[VectorStoreIndex] = None
        self._index_loaded_at: Optional[float] = None

    @property
    def index_dir(self) -> str:
        """获取索引目录路径"""
        return self.storage_dir

    def index_exists(self) -> bool:
        """
        检查索引是否存在

        Returns:
            bool: 索引是否存在
        """
        if not os.path.exists(self.storage_dir):
            return False

        # 检查必要的索引文件
        required_files = ["docstore.json", "index_store.json"]
        for file in required_files:
            if not self.file_exists(file):
                return False

        return True

    def save_index(self, index: VectorStoreIndex) -> None:
        """
        保存索引到磁盘

        Args:
            index: 要保存的索引
        """
        if not index:
            raise ValueError("索引不能为空")

        print(f"正在保存索引到: {self.storage_dir}")
        start_time = time.time()

        # 确保目录存在
        self.ensure_directory()

        # 创建/获取存储上下文并持久化
        try:
            if getattr(index, "storage_context", None) is not None:
                index.storage_context.persist(persist_dir=self.storage_dir)
            else:
                storage_context = StorageContext.from_defaults()
                storage_context.persist(persist_dir=self.storage_dir)
        except Exception as e:
            raise RuntimeError(f"索引持久化失败: {e}")

        save_time = time.time() - start_time
        self._index = index
        self._index_loaded_at = time.time()

        print(f"索引保存完成，耗时: {save_time:.2f}秒")

    def load_index(self) -> Optional[VectorStoreIndex]:
        """
        从磁盘加载索引

        Returns:
            Optional[VectorStoreIndex]: 加载的索引，如果失败则返回None
        """
        if not self.index_exists():
            print(f"索引不存在: {self.storage_dir}")
            return None

        if self._index is not None:
            print("索引已加载在内存中")
            return self._index

        print(f"正在从 {self.storage_dir} 加载索引...")
        start_time = time.time()

        try:
            # 创建存储上下文并加载索引
            storage_context = StorageContext.from_defaults(
                persist_dir=self.storage_dir
            )
            self._index = load_index_from_storage(storage_context)
            self._index_loaded_at = time.time()

            load_time = time.time() - start_time
            print(f"索引加载完成，耗时: {load_time:.2f}秒")

            return self._index

        except Exception as e:
            print(f"索引加载失败: {e}")
            return None

    def create_index(self, documents: List[Document]) -> VectorStoreIndex:
        """
        创建新的索引

        Args:
            documents: 文档列表

        Returns:
            VectorStoreIndex: 创建的索引
        """
        if not documents:
            raise ValueError("文档列表不能为空")

        print(f"正在创建索引，文档数量: {len(documents)}")
        start_time = time.time()

        # 创建索引并绑定新的存储上下文（不预填充持久化目录，避免空目录触发文件不存在）
        storage_context = StorageContext.from_defaults()
        index = VectorStoreIndex.from_documents(
            documents,
            storage_context=storage_context,
        )
        creation_time = time.time() - start_time

        print(f"索引创建完成，耗时: {creation_time:.2f}秒")
        self._index = index
        self._index_loaded_at = time.time()

        # 自动保存索引
        self.save_index(index)

        return index

    def update_index(self, documents: List[Document]) -> VectorStoreIndex:
        """
        增量更新索引

        Args:
            documents: 新增文档列表

        Returns:
            VectorStoreIndex: 更新后的索引
        """
        if not documents:
            raise ValueError("增量文档列表不能为空")

        index = self.load_index()
        if index is None:
            # 如果不存在索引，则创建新索引
            return self.create_index(documents)

        print(f"增量更新索引，新增文档: {len(documents)}")
        start_time = time.time()

        try:
            if hasattr(index, "insert"):
                index.insert(documents)
            elif hasattr(index, "insert_nodes"):
                index.insert_nodes(documents)
            else:
                raise RuntimeError("当前索引实例不支持增量插入")
            elapsed = time.time() - start_time
            print(f"增量更新完成，耗时: {elapsed:.2f}秒")
        except Exception as e:
            raise RuntimeError(f"增量更新失败: {e}")

        # 持久化
        self.save_index(index)
        return index

    def get_or_create_index(self, documents: List[Document]) -> VectorStoreIndex:
        """
        获取现有索引或创建新索引

        Args:
            documents: 文档列表

        Returns:
            VectorStoreIndex: 索引实例
        """
        if self.index_exists():
            print("发现已存在的索引，正在加载...")
            return self.load_index()
        else:
            print("未发现现有索引，正在创建...")
            return self.create_index(documents)

    def get_index_info(self) -> Dict[str, Any]:
        """
        获取索引信息

        Returns:
            Dict[str, Any]: 索引信息
        """
        info = {
            "storage_dir": self.storage_dir,
            "index_exists": self.index_exists(),
            "index_loaded": self._index is not None,
            "loaded_at": self._index_loaded_at,
            "storage_files": []
        }

        # 获取存储文件信息
        if os.path.exists(self.storage_dir):
            for file in os.listdir(self.storage_dir):
                file_path = os.path.join(self.storage_dir, file)
                if os.path.isfile(file_path):
                    stat = os.stat(file_path)
                    info["storage_files"].append({
                        "name": file,
                        "size": stat.st_size,
                        "modified": stat.st_mtime
                    })

        # 计算总存储大小
        total_size = sum(f["size"] for f in info["storage_files"])
        info["total_storage_size"] = total_size
        info["total_storage_size_mb"] = round(total_size / (1024 * 1024), 2)

        return info

    def delete_index(self) -> bool:
        """
        删除索引

        Returns:
            bool: 是否删除成功
        """
        try:
            if os.path.exists(self.storage_dir):
                import shutil
                shutil.rmtree(self.storage_dir)
                print(f"索引目录已删除: {self.storage_dir}")

            # 清空内存中的索引
            self._index = None
            self._index_loaded_at = None

            return True
        except Exception as e:
            print(f"删除索引失败: {e}")
            return False

    def backup_index(self, backup_path: str) -> bool:
        """
        备份索引

        Args:
            backup_path: 备份路径

        Returns:
            bool: 是否备份成功
        """
        try:
            if not self.index_exists():
                print("索引不存在，无法备份")
                return False

            import shutil
            backup_dir = os.path.join(backup_path, f"index_backup_{int(time.time())}")
            shutil.copytree(self.storage_dir, backup_dir)
            print(f"索引备份完成: {backup_dir}")
            return True
        except Exception as e:
            print(f"索引备份失败: {e}")
            return False

    @property
    def index(self) -> Optional[VectorStoreIndex]:
        """获取当前索引实例"""
        return self._index

    def reload_index(self) -> Optional[VectorStoreIndex]:
        """重新加载索引"""
        self._index = None
        self._index_loaded_at = None
        return self.load_index()

    def __str__(self) -> str:
        """返回存储信息的字符串表示"""
        info = self.get_index_info()
        return f"""
=== 索引存储管理器 ===
存储目录: {info['storage_dir']}
索引存在: {'是' if info['index_exists'] else '否'}
索引已加载: {'是' if info['index_loaded'] else '否'}
加载时间: {info['loaded_at'] or 'N/A'}
存储文件数: {len(info['storage_files'])}
总存储大小: {info['total_storage_size_mb']} MB
        """.strip()